import random
import os
import torch
from .my_tools import tokenizer


device = 'cuda' if torch.cuda.is_available() else 'cpu'
# print('--- device:', device)

output_dir = os.path.join('tempdir', 'models')
os.makedirs(output_dir, exist_ok=True)
# print('--- output_dir:', output_dir)


class ModelGEN(torch.nn.Module):

    def __init__(self):
        super().__init__()
        from transformers import GPT2Config, GPT2Model

        self.config = GPT2Config(bos_token_id=tokenizer.encoder['S'],
                                 eos_token_id=tokenizer.encoder['E'],
                                 n_embd=64,
                                 n_head=4,
                                 n_layer=4,
                                 n_positions=128,
                                 vocab_size=len(tokenizer.decoder))

        self.feature = GPT2Model(self.config)

        self.fc_out = torch.nn.Linear(64, self.config.vocab_size, bias=False)

        self.to(device)
        self.train()

    def forward(self, input_ids, attention_mask):
        out = self.feature(input_ids=input_ids,
                           attention_mask=attention_mask).last_hidden_state

        return self.fc_out(out)


class ModelCLS(torch.nn.Module):

    def __init__(self):
        super().__init__()
        from transformers import BertConfig, BertModel

        self.config = BertConfig(hidden_size=64,
                                 intermediate_size=64,
                                 max_position_embeddings=128,
                                 num_attention_heads=4,
                                 num_hidden_layers=4,
                                 vocab_size=len(tokenizer.decoder))

        self.feature = BertModel(self.config)

        self.fc_out = torch.nn.Sequential(torch.nn.Dropout(p=0.1),
                                          torch.nn.Linear(64, 4),
                                          )

        from tools import get_tokenizer, get_pretrained_model__bert

        self.to(device)
        self.train()

    def forward(self, input_ids, attention_mask):
        out = self.feature(input_ids=input_ids,
                           attention_mask=attention_mask).pooler_output

        return self.fc_out(out)


class ModelPPO(torch.nn.Module):

    def __init__(self, model_gen):
        super().__init__()
        self.model_gen = model_gen
        self.v_head = torch.nn.Sequential(torch.nn.Dropout(0.1),
                                          torch.nn.Linear(64, 1))

        self.to(device)
        self.train()

    def forward(self, input_ids, attention_mask):
        last_hidden_state = self.model_gen.feature(
            input_ids=input_ids,
            attention_mask=attention_mask,
            output_hidden_states=True).last_hidden_state

        logits = self.model_gen.fc_out(last_hidden_state)
        value = self.v_head(last_hidden_state).squeeze(-1)

        return logits, value


if __name__ == '__main__':
    pass




